# Copyright 2023 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import annotations

import argparse
import dataclasses
import datetime as dt
import enum
from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Self, \
    TypeAlias

from typing_extensions import override

from crossbench.config import ConfigObject, ConfigParser
from crossbench.parse import DurationParser, NumberParser, ObjectParser
from crossbench.str_enum_with_help import StrEnumWithHelp

if TYPE_CHECKING:
  Number: TypeAlias = float | int


@enum.unique
class ValidationMode(StrEnumWithHelp):
  THROW = ("throw", "Strict mode, throw and abort on env issues")
  PROMPT = ("prompt", "Prompt to accept potential env issues")
  WARN = ("warn", "Only display a warning for env issue")
  SKIP = ("skip", "Don't perform any env validation")


def merge_bool(name: str, left: Optional[bool],
               right: Optional[bool]) -> Optional[bool]:
  if left is None:
    return right
  if right is None:
    return left
  if left != right:
    raise ValueError(f"Conflicting merge values for {name}: "
                     f"{left} vs. {right}")
  return left


def merge_number_max(name: str, left: Optional[Number],
                     right: Optional[Number]) -> Optional[Number]:
  del name
  if left is None:
    return right
  if right is None:
    return left
  return max(left, right)


def merge_number_min(name: str, left: Optional[Number],
                     right: Optional[Number]) -> Optional[Number]:
  del name
  if left is None:
    return right
  if right is None:
    return left
  return min(left, right)


def merge_str_list(name: str, left: Optional[list[str]],
                   right: Optional[list[str]]) -> Optional[list[str]]:
  del name
  if left is None:
    return right
  if right is None:
    return left
  return left + right


def merge_duration_max(name: str, left: Optional[dt.timedelta],
                       right: Optional[dt.timedelta]) -> Optional[dt.timedelta]:
  del name
  if not left:
    return right
  if not right:
    return left
  return max(left, right)


ENV_CONFIG_PRESETS: dict[str, "EnvConfig"] = {}


@dataclasses.dataclass(frozen=True)
class EnvConfig(ConfigObject):
  IGNORE: ClassVar[None] = None

  browser_allow_background: bool | None = IGNORE
  browser_allow_existing_process: bool | None = IGNORE
  browser_is_headless: bool | None = IGNORE
  cpu_max_usage_percent: float | None = IGNORE
  cpu_min_relative_speed: float | None = IGNORE
  disk_min_free_space_gib: float | None = IGNORE
  power_use_battery: bool | None = IGNORE
  require_probes: bool | None = IGNORE
  screen_allow_autobrightness: bool | None = IGNORE
  screen_brightness_percent: int | None = IGNORE
  screen_refresh_rate: int | None = IGNORE
  system_allow_monitoring: bool | None = IGNORE
  system_forbidden_process_names: list[str] | None = IGNORE
  system_min_uptime: dt.timedelta | None = IGNORE

  @classmethod
  def default(cls) -> EnvConfig:
    return ENV_CONFIG_PRESETS["default"]

  @classmethod
  @override
  def parse_str(cls, value: str) -> EnvConfig:
    value = ObjectParser.non_empty_str(value)
    if preset := ENV_CONFIG_PRESETS.get(value):
      return preset
    raise argparse.ArgumentTypeError(
        f"Unknown host config preset {repr(value)}. "
        f"Choices are {','.join(ENV_CONFIG_PRESETS.keys())}")

  @classmethod
  @override
  def parse_dict(cls, config: dict[str, Any], **kwargs) -> Self:
    if "env" in config:
      config = config["env"]
    return super().parse_dict(config, **kwargs)

  @classmethod
  @override
  def config_parser(cls) -> ConfigParser[Self]:
    parser = ConfigParser(cls)
    parser.add_argument("browser_allow_background", type=ObjectParser.bool)
    parser.add_argument(
        "browser_allow_existing_process",
        type=ObjectParser.bool,
        default=cls.IGNORE)
    parser.add_argument("browser_is_headless", type=ObjectParser.bool)
    parser.add_argument(
        "cpu_max_usage_percent",
        type=NumberParser.int_range(0, 100),
        default=cls.IGNORE)
    parser.add_argument(
        "cpu_min_relative_speed",
        type=NumberParser.int_range(0, 1),
        default=cls.IGNORE)
    parser.add_argument(
        "disk_min_free_space_gib",
        type=NumberParser.positive_float,
        default=cls.IGNORE)
    parser.add_argument("power_use_battery", type=ObjectParser.bool)
    parser.add_argument("require_probes", type=ObjectParser.bool)
    parser.add_argument(
        "screen_allow_autobrightness",
        type=ObjectParser.bool,
        default=cls.IGNORE)
    parser.add_argument("screen_brightness_percent", type=int)
    parser.add_argument(
        "screen_refresh_rate",
        type=NumberParser.int_range(30, 240),
        default=cls.IGNORE)
    parser.add_argument("system_allow_monitoring", type=ObjectParser.bool)
    parser.add_argument(
        "system_forbidden_process_names", type=str, is_list=True)
    parser.add_argument(
        "system_min_uptime", type=DurationParser.positive_or_zero_duration)
    return parser

  def merge(self, other: EnvConfig) -> EnvConfig:
    mergers: dict[str, Callable[[str, Any, Any], Any]] = {
        "browser_allow_background": merge_bool,
        "browser_allow_existing_process": merge_bool,
        "browser_is_headless": merge_bool,
        "cpu_max_usage_percent": merge_number_min,
        "cpu_min_relative_speed": merge_number_max,
        "disk_min_free_space_gib": merge_number_max,
        "power_use_battery": merge_bool,
        "require_probes": merge_bool,
        "screen_allow_autobrightness": merge_bool,
        "screen_brightness_percent": merge_number_max,
        "screen_refresh_rate": merge_number_max,
        "system_allow_monitoring": merge_bool,
        "system_forbidden_process_names": merge_str_list,
        "system_min_uptime": merge_duration_max
    }
    kwargs = {}
    for name, merger in mergers.items():
      self_value = getattr(self, name)
      other_value = getattr(other, name)
      kwargs[name] = merger(name, self_value, other_value)
    return EnvConfig(**kwargs)


_config_default = EnvConfig()
_config_strict = EnvConfig(
    cpu_max_usage_percent=98,
    cpu_min_relative_speed=1,
    system_allow_monitoring=False,
    browser_allow_existing_process=False,
    require_probes=True,
    system_min_uptime=dt.timedelta(minutes=5))
_config_battery: EnvConfig = _config_strict.merge(
    EnvConfig(power_use_battery=True))
_config_power: EnvConfig = _config_strict.merge(
    EnvConfig(power_use_battery=False))
_config_catan: EnvConfig = _config_strict.merge(
    EnvConfig(
        screen_brightness_percent=65,
        system_forbidden_process_names=["terminal", "iterm2"],
        screen_allow_autobrightness=False))

ENV_CONFIG_PRESETS.update({
    "default": _config_default,
    "strict": _config_strict,
    "battery": _config_battery,
    "power": _config_power,
    "catan": _config_catan,
})
